/*
 * Copyright (C) 2025 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use crate::common::constants::*;
use crate::common::types::*;
use crate::log_e;
use crate::traits::crypto_engine::CryptoEngineRegistry;
use crate::traits::time_keeper::TimeKeeperRegistry;

const SESSION_VALIDITY_PERIOD_MS: u64 = 10 * 60 * 1000;

#[derive(Clone)]
pub struct EditSession {
    user_id: i32,
    open_time: u64,
    auth_token_valid_time: u64,
    challenge: Challenge,
    root_secret: Option<RootSecret>,
}

impl EditSession {
    pub fn new(user_id: i32) -> Result<Self, ErrorCode> {
        let mut challenge = Challenge::default();
        CryptoEngineRegistry::get().secure_random(&mut challenge)?;
        let now = TimeKeeperRegistry::get().get_system_time()?;
        Ok(Self {
            user_id: user_id,
            open_time: now,
            auth_token_valid_time: now,
            challenge: challenge,
            root_secret: None,
        })
    }

    pub fn check_valid(&mut self, user_id: Option<i32>) -> Result<(), ErrorCode> {
        let now = TimeKeeperRegistry::get().get_system_time()?;
        if now < self.open_time {
            log_e!(
                "invalid open time, now: {}, open_time: {}",
                now,
                self.open_time
            );
            return Err(ErrorCode::GeneralError);
        }

        if now
            > self
                .open_time
                .checked_add(SESSION_VALIDITY_PERIOD_MS)
                .unwrap_or(u64::MAX)
        {
            log_e!(
                "session expired, now: {}, open_time: {}",
                now,
                self.open_time
            );
            self.root_secret = None;
            return Err(ErrorCode::GeneralError);
        }

        if let Some(expected_user_id) = user_id {
            if self.user_id != expected_user_id {
                log_e!(
                    "user id mismatch, expected: {}, actual: {}",
                    expected_user_id,
                    self.user_id
                );
                return Err(ErrorCode::GeneralError);
            }
        }
        Ok(())
    }

    pub fn challenge(&self) -> &Challenge {
        &self.challenge
    }

    pub fn user_id(&self) -> UserId {
        self.user_id
    }

    fn set_root_secret(&mut self, root_secret: RootSecret) -> () {
        self.root_secret = Some(root_secret);
    }

    pub fn get_root_secret(&self) -> Result<RootSecret, ErrorCode> {
        match self.root_secret {
            Some(ref root_secret) => Ok(*root_secret),
            None => {
                log_e!("root secret not set");
                Err(ErrorCode::GeneralError)
            }
        }
    }

    fn refresh_token_valid_time(&mut self) -> () {
        self.auth_token_valid_time = match TimeKeeperRegistry::get().get_system_time() {
            Ok(time) => time,
            Err(_) => {
                log_e!("get system time failed");
                return;
            }
        };
    }

    pub fn is_token_time_valid(&self, token_time: u64) -> bool {
        token_time >= self.auth_token_valid_time
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ut_registry_guard;
    use crate::traits::crypto_engine::{AesGcmResult, KeyPair};
    use crate::utils::scope_guard::ScopeGuard;
    use crate::Vec;
    use crate::traits::time_keeper::MockTimeKeeper;
    use crate::traits::crypto_engine::MockCryptoEngine;

    #[test]
    fn test_new_success() {
        let _guard = ut_registry_guard!();

        let mut mock_time_keeper = MockTimeKeeper::new();
        mock_time_keeper
            .expect_get_system_time()
            .returning(move || Ok(1000));
        TimeKeeperRegistry::set(Box::new(mock_time_keeper));

        let challenge: Challenge = std::array::from_fn(|i| (i + 1) as u8);
        let mut mock_crypto_engine = MockCryptoEngine::new();
        mock_crypto_engine
            .expect_secure_random()
            .returning(move |buffer| {
                buffer.copy_from_slice(&challenge);
                Ok(())
            });
        CryptoEngineRegistry::set(Box::new(mock_crypto_engine));

        let session = EditSession::new(42);
        assert!(session.is_ok());
        let session = session.unwrap();
        assert_eq!(session.user_id, 42);
        assert_eq!(session.open_time, 1000);
        assert_eq!(session.auth_token_valid_time, 1000);
        assert_eq!(session.challenge, challenge);
        assert!(session.root_secret.is_none());
    }

    #[test]
    fn test_new_crypto_error() {
        let _guard = ut_registry_guard!();

        let mut mock_time_keeper = MockTimeKeeper::new();
        mock_time_keeper
            .expect_get_system_time()
            .returning(move || Ok(1000));
        TimeKeeperRegistry::set(Box::new(mock_time_keeper));

        let mut mock_crypto_engine = MockCryptoEngine::new();
        mock_crypto_engine
            .expect_secure_random()
            .returning(move |_| Err(ErrorCode::GeneralError));
        CryptoEngineRegistry::set(Box::new(mock_crypto_engine));

        let session = EditSession::new(42);
        assert!(session.is_err());
        if let Err(e) = session {
            assert_eq!(e, ErrorCode::GeneralError);
        }
    }

    #[test]
    fn test_new_time_error() {
        let _guard = ut_registry_guard!();

        let mut mock_time_keeper = MockTimeKeeper::new();
        mock_time_keeper
            .expect_get_system_time()
            .returning(move || Err(ErrorCode::GeneralError));
        TimeKeeperRegistry::set(Box::new(mock_time_keeper));

        let mut mock_crypto_engine = MockCryptoEngine::new();
        mock_crypto_engine
            .expect_secure_random()
            .returning(move |_| Ok(()));
        CryptoEngineRegistry::set(Box::new(mock_crypto_engine));

        let session = EditSession::new(42);
        assert!(session.is_err());
        if let Err(e) = session {
            assert_eq!(e, ErrorCode::GeneralError);
        }
    }

    #[test]
    fn test_check_valid_success() {
        let _guard = ut_registry_guard!();

        let mut mock_time_keeper = MockTimeKeeper::new();
        mock_time_keeper
            .expect_get_system_time()
            .returning(move || Ok(1000));
        TimeKeeperRegistry::set(Box::new(mock_time_keeper));

        let mut mock_crypto_engine = MockCryptoEngine::new();
        mock_crypto_engine
            .expect_secure_random()
            .returning(move |_| Ok(()));
        CryptoEngineRegistry::set(Box::new(mock_crypto_engine));

        let mut session = EditSession::new(42).unwrap();

        // Set time to before expiry (within validity period)
        let mut mock_time_keeper = MockTimeKeeper::new();
        mock_time_keeper
            .expect_get_system_time()
            .returning(move || Ok(1000 + SESSION_VALIDITY_PERIOD_MS - 1));
        TimeKeeperRegistry::set(Box::new(mock_time_keeper));

        let result = session.check_valid(Some(42));
        assert!(result.is_ok());
    }

    #[test]
    fn test_check_valid_expired() {
        let _guard = ut_registry_guard!();

        let mut mock_time_keeper = MockTimeKeeper::new();
        mock_time_keeper
            .expect_get_system_time()
            .returning(move || Ok(1000));
        TimeKeeperRegistry::set(Box::new(mock_time_keeper));

        let mut mock_crypto_engine = MockCryptoEngine::new();
        mock_crypto_engine
            .expect_secure_random()
            .returning(move |_| Ok(()));
        CryptoEngineRegistry::set(Box::new(mock_crypto_engine));

        let mut session = EditSession::new(42).unwrap();
        session.root_secret = Some(RootSecret::default());

        // Set time to after expiry (exceeds validity period)
        let mut mock_time_keeper = MockTimeKeeper::new();
        mock_time_keeper
            .expect_get_system_time()
            .returning(move || Ok(1000 + SESSION_VALIDITY_PERIOD_MS + 1));
        TimeKeeperRegistry::set(Box::new(mock_time_keeper));

        let result = session.check_valid(Some(42));
        assert!(result.is_err());
        if let Err(e) = result {
            assert_eq!(e, ErrorCode::GeneralError);
        }
        assert!(session.root_secret.is_none());
    }

    #[test]
    fn test_check_valid_user_mismatch() {
        let _guard = ut_registry_guard!();

        let mut mock_time_keeper = MockTimeKeeper::new();
        mock_time_keeper
            .expect_get_system_time()
            .returning(move || Ok(1000));
        TimeKeeperRegistry::set(Box::new(mock_time_keeper));

        let mut mock_crypto_engine = MockCryptoEngine::new();
        mock_crypto_engine
            .expect_secure_random()
            .returning(move |_| Ok(()));
        CryptoEngineRegistry::set(Box::new(mock_crypto_engine));

        let mut session = EditSession::new(42).unwrap();

        // Set time to within validity period
        let mut mock_time_keeper = MockTimeKeeper::new();
        mock_time_keeper
            .expect_get_system_time()
            .returning(move || Ok(1000 + SESSION_VALIDITY_PERIOD_MS - 1));
        TimeKeeperRegistry::set(Box::new(mock_time_keeper));

        let result = session.check_valid(Some(43));
        assert!(result.is_err());
        if let Err(e) = result {
            assert_eq!(e, ErrorCode::GeneralError);
        }
    }

    #[test]
    fn test_check_valid_no_user_check() {
        let _guard = ut_registry_guard!();

        let mut mock_time_keeper = MockTimeKeeper::new();
        mock_time_keeper
            .expect_get_system_time()
            .returning(move || Ok(1000));
        TimeKeeperRegistry::set(Box::new(mock_time_keeper));

        let mut mock_crypto_engine = MockCryptoEngine::new();
        mock_crypto_engine
            .expect_secure_random()
            .returning(move |_| Ok(()));
        CryptoEngineRegistry::set(Box::new(mock_crypto_engine));

        let mut session = EditSession::new(42).unwrap();

        // Set time to within validity period
        let mut mock_time_keeper = MockTimeKeeper::new();
        mock_time_keeper
            .expect_get_system_time()
            .returning(move || Ok(1000 + SESSION_VALIDITY_PERIOD_MS - 1));
        TimeKeeperRegistry::set(Box::new(mock_time_keeper));

        let result = session.check_valid(None);
        assert!(result.is_ok());
    }

    #[test]
    fn test_challenge() {
        let _guard = ut_registry_guard!();

        let expected_challenge: Challenge = std::array::from_fn(|i| i as u8);

        let mut mock_time_keeper = MockTimeKeeper::new();
        mock_time_keeper
            .expect_get_system_time()
            .returning(move || Ok(1000));
        TimeKeeperRegistry::set(Box::new(mock_time_keeper));

        let mut mock_crypto_engine = MockCryptoEngine::new();
        mock_crypto_engine
            .expect_secure_random()
            .returning(move |buffer| {
                buffer.copy_from_slice(&expected_challenge);
                Ok(())
            });
        CryptoEngineRegistry::set(Box::new(mock_crypto_engine));

        let session = EditSession::new(42).unwrap();
        let challenge = session.challenge();
        assert_eq!(challenge.len(), 32);
        // Check that the mock data was written
        for i in 0..32 {
            assert_eq!(challenge[i], i as u8);
        }
    }

    #[test]
    fn test_user_id() {
        let _guard = ut_registry_guard!();

        let mut mock_time_keeper = MockTimeKeeper::new();
        mock_time_keeper
            .expect_get_system_time()
            .returning(move || Ok(1000));
        TimeKeeperRegistry::set(Box::new(mock_time_keeper));

        let mut mock_crypto_engine = MockCryptoEngine::new();
        mock_crypto_engine
            .expect_secure_random()
            .returning(move |_| Ok(()));
        CryptoEngineRegistry::set(Box::new(mock_crypto_engine));

        let session = EditSession::new(42).unwrap();
        assert_eq!(session.user_id(), 42);
    }

    #[test]
    fn test_set_root_secret() {
        let _guard = ut_registry_guard!();

        let mut mock_time_keeper = MockTimeKeeper::new();
        mock_time_keeper
            .expect_get_system_time()
            .returning(move || Ok(1000));
        TimeKeeperRegistry::set(Box::new(mock_time_keeper));

        let mut mock_crypto_engine = MockCryptoEngine::new();
        mock_crypto_engine
            .expect_secure_random()
            .returning(move |_| Ok(()));
        CryptoEngineRegistry::set(Box::new(mock_crypto_engine));

        let mut session = EditSession::new(42).unwrap();
        assert!(session.root_secret.is_none());

        let secret = RootSecret::default();
        session.set_root_secret(secret);
        assert!(session.root_secret.is_some());
        // Check the content without consuming the Option
        assert_eq!(session.root_secret.as_ref().unwrap(), &secret);
    }

    #[test]
    fn test_get_root_secret_success() {
        let _guard = ut_registry_guard!();

        let mut mock_time_keeper = MockTimeKeeper::new();
        mock_time_keeper
            .expect_get_system_time()
            .returning(move || Ok(1000));
        TimeKeeperRegistry::set(Box::new(mock_time_keeper));

        let mut mock_crypto_engine = MockCryptoEngine::new();
        mock_crypto_engine
            .expect_secure_random()
            .returning(move |_| Ok(()));
        CryptoEngineRegistry::set(Box::new(mock_crypto_engine));

        let mut session = EditSession::new(42).unwrap();
        let secret = RootSecret::default();
        session.root_secret = Some(secret);

        let retrieved = session.get_root_secret();
        assert!(retrieved.is_ok());
        assert_eq!(retrieved.unwrap(), secret);
    }

    #[test]
    fn test_get_root_secret_none() {
        let _guard = ut_registry_guard!();

        let mut mock_time_keeper = MockTimeKeeper::new();
        mock_time_keeper
            .expect_get_system_time()
            .returning(move || Ok(1000));
        TimeKeeperRegistry::set(Box::new(mock_time_keeper));

        let mut mock_crypto_engine = MockCryptoEngine::new();
        mock_crypto_engine
            .expect_secure_random()
            .returning(move |_| Ok(()));
        CryptoEngineRegistry::set(Box::new(mock_crypto_engine));

        let session = EditSession::new(42).unwrap();
        let retrieved = session.get_root_secret();
        assert!(retrieved.is_err());
        if let Err(e) = retrieved {
            assert_eq!(e, ErrorCode::GeneralError);
        }
    }

    #[test]
    fn test_refresh_token_valid_time_success() {
        let _guard = ut_registry_guard!();

        let mut mock_time_keeper = MockTimeKeeper::new();
        mock_time_keeper
            .expect_get_system_time()
            .returning(move || Ok(1000));
        TimeKeeperRegistry::set(Box::new(mock_time_keeper));

        let mut mock_crypto_engine = MockCryptoEngine::new();
        mock_crypto_engine
            .expect_secure_random()
            .returning(move |_| Ok(()));
        CryptoEngineRegistry::set(Box::new(mock_crypto_engine));

        let mut session = EditSession::new(42).unwrap();
        assert_eq!(session.auth_token_valid_time, 1000);

        let mut mock_time_keeper = MockTimeKeeper::new();
        mock_time_keeper
            .expect_get_system_time()
            .returning(move || Ok(2000));
        TimeKeeperRegistry::set(Box::new(mock_time_keeper));

        session.refresh_token_valid_time();
        assert_eq!(session.auth_token_valid_time, 2000);
    }

    #[test]
    fn test_refresh_token_valid_time_error() {
        let _guard = ut_registry_guard!();

        let mut mock_time_keeper = MockTimeKeeper::new();
        mock_time_keeper
            .expect_get_system_time()
            .returning(move || Ok(1000));
        TimeKeeperRegistry::set(Box::new(mock_time_keeper));

        let mut mock_crypto_engine = MockCryptoEngine::new();
        mock_crypto_engine
            .expect_secure_random()
            .returning(move |_| Ok(()));
        CryptoEngineRegistry::set(Box::new(mock_crypto_engine));

        let mut session = EditSession::new(42).unwrap();
        assert_eq!(session.auth_token_valid_time, 1000);

        let mut mock_time_keeper = MockTimeKeeper::new();
        mock_time_keeper
            .expect_get_system_time()
            .returning(move || Err(ErrorCode::GeneralError));
        TimeKeeperRegistry::set(Box::new(mock_time_keeper));

        session.refresh_token_valid_time();
        // Should not change when time retrieval fails
        assert_eq!(session.auth_token_valid_time, 1000);
    }

    #[test]
    fn test_is_token_time_valid_true() {
        let _guard = ut_registry_guard!();

        let mut mock_time_keeper = MockTimeKeeper::new();
        mock_time_keeper
            .expect_get_system_time()
            .returning(move || Ok(1000));
        TimeKeeperRegistry::set(Box::new(mock_time_keeper));

        let mut mock_crypto_engine = MockCryptoEngine::new();
        mock_crypto_engine
            .expect_secure_random()
            .returning(move |_| Ok(()));
        CryptoEngineRegistry::set(Box::new(mock_crypto_engine));

        let mut session = EditSession::new(42).unwrap();
        session.auth_token_valid_time = 1000;

        assert!(session.is_token_time_valid(1000));
        assert!(session.is_token_time_valid(1500));
        assert!(session.is_token_time_valid(2000));
    }

    #[test]
    fn test_is_token_time_valid_false() {
        let _guard = ut_registry_guard!();

        let mut mock_time_keeper = MockTimeKeeper::new();
        mock_time_keeper
            .expect_get_system_time()
            .returning(move || Ok(1000));
        TimeKeeperRegistry::set(Box::new(mock_time_keeper));

        let mut mock_crypto_engine = MockCryptoEngine::new();
        mock_crypto_engine
            .expect_secure_random()
            .returning(move |_| Ok(()));
        CryptoEngineRegistry::set(Box::new(mock_crypto_engine));

        let mut session = EditSession::new(42).unwrap();
        session.auth_token_valid_time = 1000;

        assert!(!session.is_token_time_valid(999));
        assert!(!session.is_token_time_valid(500));
        assert!(!session.is_token_time_valid(0));
    }
}
